//////////////////////////////////////////////////////////////////////////
//
// This file is a part of UKROP project.
// Copyright (c) Vladimir Gumenuk
//
// PROJECT:         Xort Custom Debug Engine
// FILE:            Engine.cpp : Implementation of CEngine
// AUTHOR:          Vladimir Gumenuk
// DATE:            Dec-2012 -- Jan-2013
// MODIFICATIONS:
//
//////////////////////////////////////////////////////////////////////////

#include "stdafx.h"
#include "xort.h"
#include "xruntime.h"
#include "DebugProperty.h"
#include "DebugStack.h"
#include "Program.h"
#include "Process.h"
#include "Engine.h"
#include "EventBase.h"
#include "dllmain.h"

// CEngine
CEngine::CEngine()
{
    ATLTRACE( "CEngine::CEngine()\n" );
    m_wLocale       = 0;
    m_hStdInput     = 0;
    m_hStdOutput    = 0;
    m_hStdError     = 0;
    m_bTraceSettings = FALSE;
}

CEngine::~CEngine()
{
    ATLTRACE( "CEngine::~CEngine()\n" );
}

HRESULT CEngine::FinalConstruct()
{
    return S_OK;
}

void CEngine::FinalRelease()
{

}

////////////////////////////////////////////////////////////////////////////// 
//IDebugEngine2 methods 
STDMETHODIMP CEngine::EnumPrograms(IEnumDebugPrograms2** ppEnum) 
{ 
    ATLTRACE( "CEngine::EnumPrograms()\n" );
    return E_NOTIMPL;
}

STDMETHODIMP CEngine::SetException(EXCEPTION_INFO* pException) 
{
    ATLTRACE( "CEngine::SetException()\n" );
    return E_NOTIMPL;
}

STDMETHODIMP CEngine::RemoveSetException(EXCEPTION_INFO* pException) 
{
    ATLTRACE( "CEngine::RemoveSetException()\n" );
    return E_NOTIMPL;
}

STDMETHODIMP CEngine::RemoveAllSetExceptions(REFGUID guidType) 
{
    ATLTRACE( "CEngine::RemoveAllSetExceptions()\n" );
    return E_NOTIMPL;
}

STDMETHODIMP CEngine::GetEngineId(GUID *pguidEngine) 
{
    ATLTRACE( "CEngine::GetEngineId()\n" );
    *pguidEngine = __uuidof(Engine);
    return S_OK;
}

STDMETHODIMP CEngine::DestroyProgram( IDebugProgram2* pProgram )
{
    ATLTRACE( "CEngine::DestroyProgram()\n" );
    if( !pProgram )
        return E_POINTER;

    HRESULT hr = S_FALSE;
    if( m_spProgram )
    {
        m_spProgram->Terminate();

        //CProgramDestroyEvent *pProgDestroy = new CProgramDestroyEvent( 0 );
        //pProgDestroy->SendEvent( m_spCallback, this, m_spProcess, m_spProgram, NULL );

        m_spProgram.Release();

        hr = S_OK;
    }

    return hr;
}

STDMETHODIMP CEngine::SetLocale(WORD wLangID) 
{
    ATLTRACE( "CEngine::SetLocale() -> %d\n", wLangID );
    m_wLocale = wLangID;
    return S_OK;
}

STDMETHODIMP CEngine::SetRegistryRoot( LPCOLESTR pszRegistryRoot ) 
{
    ATLTRACE( "CEngine::SetRegistryRoot() -> %S\n", pszRegistryRoot );
    m_bsRegistryRoot = pszRegistryRoot;
    return S_OK;
}

STDMETHODIMP CEngine::SetMetric( LPCOLESTR pszMetric, VARIANT varValue ) 
{
    ATLTRACE( "CEngine::SetMetric() -> %S, varType: %d\n", pszMetric, varValue.vt );
    
    if( wcscmp( pszMetric, L"TraceSettings" ) == 0 )
        m_bTraceSettings = TRUE;

    return S_OK; // just return S_OK, ignore for now
}

STDMETHODIMP CEngine::CauseBreak() 
{
    ATLTRACE( "CEngine::CauseBreak()\n" );
    return E_NOTIMPL;
} 

STDMETHODIMP CEngine::Attach( IDebugProgram2 **rgpPrograms, IDebugProgramNode2 **rgpProgramNodes, 
    DWORD celtPrograms, IDebugEventCallback2 *pCallback, ATTACH_REASON dwReason ) 
{ 
    ATLTRACE( "CEngine::Attach() -> callback: 0x%08X\n", pCallback );
    for( DWORD cProgram = 0; cProgram < celtPrograms; cProgram++ ) 
    { 
        CProgram* pProgram = NULL; 
        if( rgpProgramNodes[cProgram] != NULL && 
            rgpProgramNodes[cProgram]->QueryInterface( IID_IXortProgram, (void**)&pProgram) == S_OK ) 
        { 
            pProgram->EngineAttach( rgpPrograms[cProgram], this, m_spProcess, pCallback, dwReason ); 
            pProgram->Release(); 
        } 
    } 

    return S_OK; 
} 

STDMETHODIMP CEngine::ContinueFromSynchronousEvent( IDebugEvent2* pEvent ) 
{ 
    ATLTRACE( "CEngine::ContinueFromSynchronousEvent()\n" );
    _AtlModule.PostMsgToMainThread( WM_CONTINUE_SYNC_EVENT ); // release main thread first

    CComPtr<IDebugBreakpointEvent2> spBPEvent;
    CComPtr<IDebugStepCompleteEvent2> spStepEvent;
    CComPtr<IDebugCanStopEvent2> spCanStopEvent;

    HRESULT hr = pEvent->QueryInterface( IID_IDebugBreakpointEvent2, (void**)&spBPEvent );
    if( SUCCEEDED( hr ) ) // verify that this is a breakpoint event
    {
        ATLTRACE( "CEngine::ContinueFromSynchronousEvent() -> continue program\n" );
        m_spProgram->Continue( NULL ); // and XVM will continue processing again
    }
    else if( SUCCEEDED( hr = pEvent->QueryInterface( IID_IDebugCanStopEvent2, (void**)&spCanStopEvent ) ) )
    {
        _AtlModule.getXProgramNoAddRef()->SignalEvent(); // WTF???
    }
    else if( SUCCEEDED( hr = pEvent->QueryInterface( IID_IDebugStepCompleteEvent2, (void**)&spStepEvent ) ) )
    {
        ATLTRACE( "CEngine::ContinueFromSynchronousEvent() -> IDebugStepCompleteEvent2!\n" );
    } else {
        ATLTRACE( "CEngine::ContinueFromSynchronousEvent() -> UNKNOWN EVENT!\n" );
    }
    
    return S_OK;
} 

STDMETHODIMP CEngine::CreatePendingBreakpoint( IDebugBreakpointRequest2 *pBPRequest,  IDebugPendingBreakpoint2** ppPendingBP ) 
{ 
    ATLTRACE( "CEngine::CreatePendingBreakpoint()\n" );
    if( !m_spProgram )
        return E_UNEXPECTED;
    if( !pBPRequest || !ppPendingBP )
        return E_POINTER;

    CComQIPtr<IXProgram> spProg( m_spProgram );
    return spProg->CreatePendingBreakpoint( pBPRequest, ppPendingBP );
} 

//////////////////////////////////////////////////////////// 
// IDebugEngineLaunch2 
STDMETHODIMP CEngine::LaunchSuspended( 
    __RPC__in_opt LPCOLESTR pszServer,
    __RPC__in_opt IDebugPort2 *pPort,
    __RPC__in_opt LPCOLESTR pszExe,
    __RPC__in_opt LPCOLESTR pszArgs,
    __RPC__in_opt LPCOLESTR pszDir,
    __RPC__in_opt BSTR bstrEnv,
    __RPC__in_opt LPCOLESTR pszOptions,
    LAUNCH_FLAGS dwLaunchFlags,
    DWORD hStdInput,
    DWORD hStdOutput,
    DWORD hStdError,
    __RPC__in_opt IDebugEventCallback2 *pCallback,
    __RPC__deref_out_opt IDebugProcess2 **ppProcess )
{
    ATLTRACE( "CEngine::LaunchSuspended() -> srv: %S, port: 0x%08X, \nexe: %S, \nargs: %S, \ndir: %S, \nenv: %S, opt: %S, falgs: 0x%08X, callback: 0x%08X\n", 
        pszServer, pPort, pszExe, pszArgs, pszDir, bstrEnv, pszOptions, dwLaunchFlags, pCallback );

    if( !pszArgs )
        return E_POINTER; // and what do I need to debug?
    if( m_spProcess )
        return E_UNEXPECTED; // already created?

    m_hStdInput  = hStdInput;
    m_hStdOutput = hStdOutput;
    m_hStdError  = hStdError;

    _AtlModule.SetMainThreadID();   // save main SDM thread id to post messages later
    _AtlModule.setCallback( pCallback );
    _AtlModule.setEngine( this );

    CEngineCreateEvent *pEngineCreated = new CEngineCreateEvent( this ); // tell SDM that engine is ready :)
    pEngineCreated->Send();

    HRESULT hr = S_OK;
    if( SUCCEEDED( hr = CProcess::Create( &m_spProcess ) ) )
    {
        CComQIPtr<IXProcess> spProc( m_spProcess );
        spProc->LaunchExecutable( pszExe ); // launch dummy process to make VS IDE happy :)

        _AtlModule.setProcess( m_spProcess );

        //CProcessCreateEvent *pEvent = new CProcessCreateEvent(); // tell SDM that process is ready - !NO, it hangs!
        //pEvent->SendEvent( m_spCallback, this, m_spProcess, NULL, NULL );

        if( SUCCEEDED( hr = CProgram::Create( &m_spProgram ) ) )
        {
            _AtlModule.setProgram( m_spProgram );

            CComQIPtr<IDebugThread2> spTh( m_spProgram ); // the program will be the thread as well
            _AtlModule.setThread( spTh );

            CComQIPtr<IXProgram> spXProg( m_spProgram );
            if( SUCCEEDED( hr = spXProg->Load( pszArgs ) ) ) // load program to debug
            {
                // Now get the IPortNotify2 interface so we can register a program node in CDebugEngine::ResumeProcess.
                CComPtr<IDebugDefaultPort2> spDefaultPort;
                if( SUCCEEDED( hr = pPort->QueryInterface( IID_IDebugDefaultPort2, (void**)&spDefaultPort ) ) )
                {
                    CComPtr<IDebugPortNotify2> spPortNotify;
                    if( SUCCEEDED( hr = spDefaultPort->GetPortNotify( &spPortNotify ) ) )
                    {
                        m_spPortNotify = spPortNotify; // Remember the port notify so we can use it in ResumeProcess.
                    }
                }
            }
        }
    }
    
    if( SUCCEEDED( hr ) && ppProcess ) // return process to SDM
    {
        *ppProcess = m_spProcess;
        (*ppProcess)->AddRef();
    } else {
        m_spProcess->Terminate();
        _AtlModule.ReleaseAllDebugObjects();
    }

    ATLTRACE( "CEngine::LaunchSuspended() -> LEAVE, hr: 0x%08X\n", hr );
    return hr;
}

STDMETHODIMP CEngine::ResumeProcess( __RPC__in_opt IDebugProcess2 *pProcess )
{
    ATLTRACE( "CEngine::ResumeProcess()!\n" );
    if( !pProcess )
        return E_POINTER;
    if( !m_spProcess || !m_spProgram )
        return E_UNEXPECTED;

    HRESULT hr = S_FALSE;
    if( isMyProcess( pProcess ) ) // Make a program node for this process
    {
        CComQIPtr<IDebugProgramNode2> spProgramNode( m_spProgram );
        if( SUCCEEDED( hr = m_spPortNotify->AddProgramNode( spProgramNode ) ) )
        {
            ATLTRACE( "AddProgramNode - OK!\n" );
        }
    }

    return hr; // return S_OK, our process is resumed :)
}

STDMETHODIMP CEngine::CanTerminateProcess( __RPC__in_opt IDebugProcess2 *pProcess )
{
    ATLTRACE( "CEngine::CanTerminateProcess()!\n" );
    if( !pProcess )
        return E_POINTER;
    if( !m_spProcess )
        return E_UNEXPECTED;

    HRESULT hr = ( isMyProcess( pProcess ) ) ? S_OK : S_FALSE;
    return hr; // return S_OK, our process is resumed :)
}

STDMETHODIMP CEngine::TerminateProcess( __RPC__in_opt IDebugProcess2 *pProcess )
{
    ATLTRACE( "CEngine::TerminateProcess()!\n" );
    if( !pProcess )
        return E_POINTER;
    if( !m_spProcess )
        return E_UNEXPECTED;

    HRESULT hr = E_FAIL;
    
    if( m_spProgram )
    {
        if( SUCCEEDED( hr = m_spProgram->Terminate() ) )
        {
            m_spProgram.Release(); // delete the program
        }
    }

    if( m_spProcess )
    {
        m_spProcess->Terminate();
        //CProcessDestroyEvent *pProcDestroy = new CProcessDestroyEvent(); // NO! not yet! - IDE crashes
        //pProcDestroy->SendEvent( m_spCallback, this, m_spProcess, NULL, NULL );

        //m_spProcess.Release(); // delete process, NO, not here!
    }

    _AtlModule.ReleaseAllDebugObjects(); // !!! DO not forget about this

    ATLTRACE( "CEngine::TerminateProcess() -> LEAVE, hr: 0x%08X\n", hr );
    if( SUCCEEDED( hr ) )
    {
        PostQuitMessage( 0 );
    }

    return hr; // return S_OK, our process is resumed :)
}

BOOL CEngine::isMyProcess( IDebugProcess2 *pProcess )
{
    BOOL bRet = FALSE;
    HRESULT hr;
    GUID vmGUID;
    if( SUCCEEDED( hr = pProcess->GetProcessId( &vmGUID ) ) )
    {
        if( IsEqualGUID( IID_XortVirtualProcess, vmGUID ) )
        {
            bRet = TRUE;
        }
    }
    return bRet;
}

